import { _decorator, CCBoolean, log, math, Node, Quat, RigidBody, v3, Vec3 } from 'cc';
import { Agent } from '../../../../Runtime/Agent';
import { Academy } from '../../../../Runtime/Academy';
import { EnvironmentParameters } from '../../../../Runtime/EnvironmentParameters';
import { VectorSensor } from '../../../../Runtime/Sensors/VectorSensor';
import { ActionBuffers } from '../../../../Runtime/Actuators/IActionReceiver';
import { Axis, InputWrap } from '../../../Common/Scripts/InputWrap';
const { ccclass, property } = _decorator;

@ccclass('Ball3DAgent')
export class Ball3DAgent extends Agent {
	@property(Node)
	ball: Node = null!;

	@property(CCBoolean)
	userVecObs: boolean = true;

	private _ballRb: RigidBody = null!;
	private _resetParams: EnvironmentParameters = null!;

	private _tempVec3: Vec3 = new Vec3();
	private _tempVelocityVec3: Vec3 = new Vec3();

	inputWrap: InputWrap = null!;

	onLoad(): void {
		this._ballRb = this.ball.getComponent(RigidBody)!;
		this._resetParams = Academy.instance.environmentParameters;
		this.inputWrap = this.node.getComponent(InputWrap)!;
	}

	start() {}

	update(deltaTime: number) {}

	override collectObservations(sensor: VectorSensor): void {
		if (this.userVecObs) {
			sensor.addObservation(this.node.rotation.z);
			sensor.addObservation(this.node.rotation.x);
			this._tempVec3.set(
				this.ball.position.x - this.node.position.x,
				this.ball.position.y - this.node.position.y,
				this.ball.position.z - this.node.position.z
			);
			sensor.addObservation(this._tempVec3);
			this._ballRb.getLinearVelocity(this._tempVelocityVec3);
			sensor.addObservation(this._tempVelocityVec3);
		}
	}

	override onActionReceived(actionBuffers: ActionBuffers): void {
		const actionZ = 2 * math.clamp(actionBuffers.continuousActions.getValue(0), -1, 1);
		const actionX = 2 * math.clamp(actionBuffers.continuousActions.getValue(1), -1, 1);

		if ((this.node.rotation.z < 0.25 && actionZ > 0) || (this.node.rotation.z > -0.25 && actionZ < 0)) {
			this.node.setRotationFromEuler(0, 0, actionZ);
		}

		if ((this.node.rotation.x < 0.25 && actionX > 0) || (this.node.rotation.x > -0.25 && actionX < 0)) {
			this.node.setRotationFromEuler(actionX, 0, 0);
		}

		if (
			this.ball.position.y - this.node.position.y < -2 ||
			Math.abs(this.ball.position.x - this.node.position.x) > 3 ||
			Math.abs(this.ball.position.z - this.node.position.z) > 3
		) {
			this.setReward(-1.0);
			this.endEpisode();
		} else {
			this.setReward(0.1);
		}
	}

	override onEpisodeBegin(): void {
		this.node.setRotationFromEuler(0, math.randomRange(-30, 10), math.randomRange(-30, 10));
		this._ballRb.setLinearVelocity(v3(0, 0, 0));
		this.ball.setPosition(
			v3(
				math.randomRange(-0.5, 0.5) + this.node.position.x,
				4 + math.randomRange(-0.5, 0.5) + this.node.position.y,
				math.randomRange(-0.5, 0.5) + this.node.position.z
			)
		);

		this.setResetParameters();
	}

	override heuristic(actionsOut: ActionBuffers): void {
		const continuousActionsOut = actionsOut.continuousActions;

		continuousActionsOut.setValue(0, this.inputWrap.getAxis(Axis.Horizontal));
		continuousActionsOut.setValue(1, this.inputWrap.getAxis(Axis.Vertical));
	}

	setBall() {
		this._ballRb.mass = this._resetParams.getWithDefault('mass', 1) * 0.2;
		// let scale = this._resetParams.getWithDefault('scale', 1)*0.5;
		let scale = 0.2;
		scale = math.randomRange(0.1, 0.5);
		log('scale', scale);

		this.ball.setScale(scale, scale, scale);
	}

	setResetParameters() {
		this.setBall();
	}
}
